from agent import Agent
from agent_manager import AgentManager

class AgentPopulation(object):
    """
    A population of agents that may be loaded into a game. Defines the classes 
    of agents in the population and the number of instances of each class.
    """
    
    def __init__(self, agent_manager=None):
        self.manager = agent_manager if agent_manager else AgentManager()
        self._data = []

    def insert_entry(self, class_name, number=1, i=-1):
        """Add agent class at end, or after specified position."""
        assert(class_name in self.manager and i < len(self._data))
        entry = {'class_name':class_name,'number':number}
        if i < 0:
            self._data.append(entry)
        else:
            self._data.insert(i+1, entry)

    def remove_entry(self, i):
        """Remove agent at specified position."""
        assert(i >= 0 and i < len(self._data))
        self._data = self._data[:i] + self._data[i+1:]
 
    def modify_number(self, i, number):
        assert(i >= 0 and i < len(self._data) and number > 0)
        self._data[i]['number'] = number

    def modify_class_name(self, i, class_name):
        assert(i >= 0 and i < len(self._data))
        self._data[i]['class_name'] = class_name
        
    def is_member(self, class_name):
        return any([x['class_name'] == class_name for x in self._data])

    def population_size(self):
        return sum([x['number'] for x in self._data])

    def available_agents(self):
        return self.manager.names()
    
    def load_agent_file(self, path):
        return self.manager.load_file(path)

    def get(self):
        result = self._data
        psum = sum([x['number'] for x in self._data]) * 1.0
        for i, x in enumerate(result):
            result[i]['class'] = self.manager[x['class_name']]
            result[i]['freq'] = result[i]['number'] / psum
        return result
    
    def make_agent_instances(self):
        instances = []
        for x in self.get():
            instances += [x['class']() for i in range(x['number'])]
        return instances